#include "Init.h"
#include "HideFolder.h"
#include "Rules.h"

NTSTATUS CheckFolderAndHideIfNeed(PDEVICE_OBJECT pDeviceObject, PUNICODE_STRING pusParentFolder, PFILE_BOTH_DIR_INFORMATION pFileBothDirInformation);

NTSTATUS IoCompletionRoutine(IN PDEVICE_OBJECT pDeviceObject,   
                             IN PIRP pIRP,   
                             IN PVOID pContext)   
{
    UNREFERENCED_PARAMETER(pDeviceObject);
    UNREFERENCED_PARAMETER(pContext);

    *pIRP->UserIosb = pIRP->IoStatus;   

    if (pIRP->UserEvent)
    {
        KeSetEvent(pIRP->UserEvent, IO_DISK_INCREMENT, 0);
    }

    IoFreeIrp(pIRP);   

    return STATUS_MORE_PROCESSING_REQUIRED;   
}

NTSTATUS CheckFolderContentAndHideIfNeed(IN PFILE_OBJECT pRootFolderFO,
                                         IN OUT PFILE_BOTH_DIR_INFORMATION pFileBothDirInfoHead,
                                         IN const unsigned int uiFileBothDirInfoFullLength,
                                         IN const BOOLEAN bIsRestartScanSet,
                                         IN const BOOLEAN bIsSingleEntrySet)
{
    PDEVICE_OBJECT pRelatedDeviceObject = IoGetRelatedDeviceObject(pRootFolderFO);
    PDEVICE_OBJECT pLowestDeviceObject = IoGetDeviceAttachmentBaseRef(pRelatedDeviceObject);
    UNICODE_STRING usParentFolder = pRootFolderFO->FileName;
    NTSTATUS ntStatus = STATUS_UNSUCCESSFUL;

    UNREFERENCED_PARAMETER(bIsRestartScanSet);

    if (pRelatedDeviceObject == pLowestDeviceObject)
    {
        ASSERT(FALSE);

        // maybe ObDereferenceObject???

        return STATUS_UNSUCCESSFUL;
    }

    if (bIsSingleEntrySet)
    {
        RULE matchedRule = {0};

        do
        {
            UNICODE_STRING usFolderName = {0};
            
            usFolderName.Length = (USHORT)pFileBothDirInfoHead->FileNameLength;
            usFolderName.MaximumLength = (USHORT)pFileBothDirInfoHead->FileNameLength;
            usFolderName.Buffer = &pFileBothDirInfoHead->FileName[0];

            ntStatus = ProcessRules(pLowestDeviceObject, &usParentFolder, &usFolderName, &matchedRule);

            switch (matchedRule.ruleCategory)
            {
                case HideFolderAction:
                {
                    unsigned int uiBufferSize = uiFileBothDirInfoFullLength;
                    PVOID pBuffer = ExAllocatePoolWithTag(PagedPool, uiBufferSize, FileHidingFilterPoolTag);

                    for (;;)
                    {
                        PDEVICE_OBJECT pRelatedDeviceObject = IoGetRelatedDeviceObject(pRootFolderFO);
                        PDEVICE_OBJECT pLowerDeviceObject = IoGetLowerDeviceObject(pRelatedDeviceObject);
                        PIRP pNewIRP = IoAllocateIrp(pLowerDeviceObject->StackSize, FALSE);
                        IO_STATUS_BLOCK ioStatus = {0};
                        PIO_STACK_LOCATION pIONewStack = IoGetNextIrpStackLocation(pNewIRP);
                        KEVENT syncEvent = {0};

                        KeInitializeEvent(&syncEvent, SynchronizationEvent, FALSE);

                        pNewIRP->Tail.Overlay.OriginalFileObject = pRootFolderFO;
                        pNewIRP->Tail.Overlay.Thread = PsGetCurrentThread();
                        pNewIRP->RequestorMode = KernelMode;
                        pNewIRP->UserEvent = &syncEvent;
                        pNewIRP->UserIosb = &ioStatus;
                        pNewIRP->UserBuffer = pBuffer;

                        SetFlag(pNewIRP->Flags, IRP_SYNCHRONOUS_API);

                        RtlZeroMemory(pIONewStack, sizeof(IO_STACK_LOCATION));

                        pIONewStack->MajorFunction = IRP_MJ_DIRECTORY_CONTROL;
                        pIONewStack->MinorFunction = IRP_MN_QUERY_DIRECTORY;
                        pIONewStack->FileObject = pRootFolderFO;

                        pIONewStack->Parameters.QueryDirectory.FileIndex = 0;
                        pIONewStack->Parameters.QueryDirectory.FileInformationClass = FileBothDirectoryInformation;
                        pIONewStack->Parameters.QueryDirectory.Length = uiBufferSize;
                        pIONewStack->Parameters.QueryDirectory.FileName = NULL;

                        SetFlag(pIONewStack->Flags, SL_RETURN_SINGLE_ENTRY);

                        IoSetCompletionRoutine(pNewIRP, IoCompletionRoutine, NULL, TRUE, TRUE, TRUE);

                        ntStatus = IoCallDriver(pLowerDeviceObject, pNewIRP);

                        if (ntStatus == STATUS_PENDING) 
                        {
                            KeWaitForSingleObject(&syncEvent, Executive, KernelMode, FALSE, NULL);

                            ntStatus = ioStatus.Status;
                        }

                        if (NT_SUCCESS(ntStatus))
                        {
                            if (uiFileBothDirInfoFullLength < ioStatus.Information)
                            {
                                ASSERT(FALSE);

                                ExFreePool(pBuffer);

                                return STATUS_SUCCESS;
                            }

                            RtlCopyMemory(pFileBothDirInfoHead, pBuffer, ioStatus.Information);
                        }

                        if (STATUS_BUFFER_OVERFLOW == ntStatus) 
                        {
                            ExFreePool(pBuffer);

                            uiBufferSize *= 2;

                            pBuffer = ExAllocatePoolWithTag(PagedPool, uiBufferSize, FileHidingFilterPoolTag);

                            continue;
                        }

                        if (STATUS_NO_MORE_FILES == ntStatus)
                        {
                            ExFreePool(pBuffer);

                            ObDereferenceObject(pLowestDeviceObject);

                            return STATUS_NO_SUCH_FILE;
                        }

                        ExFreePool(pBuffer);

                        break;
                    }

                    break;
                }
                default:
                {
                    break;
                }
            }
        } while (matchedRule.ruleCategory != AllowAction);
    }
    else
    {
        CheckFolderAndHideIfNeed(pLowestDeviceObject, &pRootFolderFO->FileName, pFileBothDirInfoHead);
    }

    ObDereferenceObject(pLowestDeviceObject);

    return STATUS_SUCCESS;
}

NTSTATUS CheckFolderAndHideIfNeed(PDEVICE_OBJECT pDeviceObject, PUNICODE_STRING pusParentFolder, PFILE_BOTH_DIR_INFORMATION pFileBothDirInformation)
{
    UNICODE_STRING usFolderName ={0};
    ULONG ulCurrentOffset = 0;
    ULONG ulPreviosOffset = 0;
    RULE Rule = {0};
    NTSTATUS ntStatus = STATUS_SUCCESS;

    do 
    {
        pFileBothDirInformation = (PFILE_BOTH_DIR_INFORMATION)(((PUCHAR)pFileBothDirInformation) + ulCurrentOffset);

        usFolderName.MaximumLength = usFolderName.Length = (USHORT)pFileBothDirInformation->FileNameLength;
        usFolderName.Buffer = &pFileBothDirInformation->FileName[0];

        ulCurrentOffset = pFileBothDirInformation->NextEntryOffset;

        ntStatus = ProcessRules(pDeviceObject, pusParentFolder, &usFolderName, &Rule);

        if (!NT_SUCCESS(ntStatus))
        {
            ASSERT(FALSE);

            return ntStatus;
        }

        switch(Rule.ruleCategory)
        {
            case HideFolderAction:
            {

                // folder for hide is first, so check next and process if next entry exist
                if(0 == ulPreviosOffset)
                {
                    if (0 != pFileBothDirInformation->NextEntryOffset)
                    {
                        PFILE_BOTH_DIR_INFORMATION pNextEntry = (PFILE_BOTH_DIR_INFORMATION)((char*)pFileBothDirInformation + pFileBothDirInformation->NextEntryOffset);
                        unsigned int uiNextOffset = pFileBothDirInformation->NextEntryOffset;

                        if (0 == pNextEntry->NextEntryOffset)
                        {
                            uiNextOffset = 0;
                        }
                        else
                        {
                            uiNextOffset += pNextEntry->NextEntryOffset;
                        }

                        RtlMoveMemory(pFileBothDirInformation, pNextEntry, (char*)pNextEntry - (char*)pFileBothDirInformation + pNextEntry->FileNameLength - sizeof(WCHAR));

                        pFileBothDirInformation->NextEntryOffset = uiNextOffset;

                        ulCurrentOffset = pFileBothDirInformation->NextEntryOffset;

                        continue;
                    }
                    else
                    {
                        ASSERT(FALSE);
                    }

                    return STATUS_SUCCESS;
                }

                pFileBothDirInformation = (PFILE_BOTH_DIR_INFORMATION)(((char*)pFileBothDirInformation) - ulPreviosOffset);

                if(ulCurrentOffset)
                {
                    pFileBothDirInformation->NextEntryOffset = pFileBothDirInformation->NextEntryOffset + ulCurrentOffset;

                    ulCurrentOffset = pFileBothDirInformation->NextEntryOffset;
                }    
                else
                {
                    pFileBothDirInformation->NextEntryOffset = 0;

                    ulCurrentOffset = 0;
                }

                break;
            }
        }

        ulPreviosOffset = ulCurrentOffset;

    } while (0 != ulCurrentOffset);

    return STATUS_SUCCESS;
}